{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulating gate noise on the Rigetti Quantum Virtual Machine\n",
    "© Copyright 2017, Rigetti Computing. \n",
    "$$\n",
    "\\newcommand{ket}[1]{\\left|{#1}\\right\\rangle}\n",
    "\\newcommand{bra}[1]{\\left\\langle {#1}\\right|}\n",
    "\\newcommand{tr}{\\mathrm{Tr}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pure states vs. mixed states\n",
    "\n",
    "\n",
    "Errors in quantum computing can introduce classical uncertainty in what the underlying state is.\n",
    "When this happens we sometimes need to consider not only wavefunctions but also probabilistic sums of\n",
    "wavefunctions when we are uncertain as to which one we have. For example, if we think that an X gate\n",
    "was accidentally applied to a qubit with a 50-50 chance then we would say that there is a 50% chance\n",
    "we have the $\\ket{0}$ state and a 50% chance that we have a $\\ket{1}$ state.\n",
    "This is called an \"impure\" or\n",
    "\"mixed\"state in that it isn't just a wavefunction (which is pure) but instead a distribution over\n",
    "wavefunctions. We describe this with something called a density matrix, which is generally an\n",
    "operator. Pure states have very simple density matrices that we can write as an outer product of a\n",
    "ket vector $\\ket{\\psi}$ with its own bra version $\\bra{\\psi}=\\ket{\\psi}^\\dagger$.\n",
    "For a pure state the density matrix is simply\n",
    "\n",
    "\n",
    "$$\n",
    "   \\rho_\\psi = \\ket{\\psi}\\bra{\\psi}.\n",
    "$$\n",
    "The expectation value of an operator for a mixed state is given by\n",
    "\n",
    "$$\n",
    "   \\langle X \\rangle_\\rho = \\tr{X \\rho}\n",
    "$$\n",
    "where $\\tr{A}$ is the trace of an operator, which is the sum of its diagonal elements\n",
    "which is independent of choice of basis.\n",
    "Pure state density matrices satisfy\n",
    "\n",
    "$$\n",
    "   \\rho \\text{ is pure } \\Leftrightarrow \\rho^2 = \\rho\n",
    "$$\n",
    "which you can easily verify for $\\rho_\\psi$ assuming that the state is normalized.\n",
    "If we want to describe a situation with classical uncertainty between states $\\rho_1$ and\n",
    "$\\rho_2$, then we can take their weighted sum\n",
    "$$\n",
    "   \\rho = p \\rho_1 + (1-p) \\rho_2\n",
    "$$\n",
    "where $p\\in [0,1]$ gives the classical probability that the state is $\\rho_1$.\n",
    "\n",
    "Note that classical uncertainty in the wavefunction is markedly different from superpositions.\n",
    "We can represent superpositions using wavefunctions, but use density matrices to describe\n",
    "distributions over wavefunctions. You can read more about density matrices [here](https://en.wikipedia.org/wiki/Density_matrix)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quantum gate errors\n",
    "## What are they?\n",
    "For a quantum gate given by its unitary operator $U$, a \"quantum gate error\" describes the scenario in which the actually induces transformation deviates from $\\ket{\\psi} \\mapsto U\\ket{\\psi}$. \n",
    "There are two basic types of quantum gate errors:\n",
    "\n",
    "1. **coherent errors** are those that preserve the purity of the input state, i.e., instead of the above mapping we carry out a perturbed, but unitary operation $\\ket{\\psi} \\mapsto \\tilde{U}\\ket{\\psi}$, where $\\tilde{U} \\neq U$.  \n",
    "2. **incoherent errors** are those that do not preserve the purity of the input state, \n",
    "    in this case we must actually represent the evolution in terms of density matrices. \n",
    "    The state $\\rho := \\ket{\\psi}\\bra{\\psi}$ is then mapped as \n",
    "    $$\n",
    "    \\rho \\mapsto \\sum_{j=1}^n K_j\\rho K_j^\\dagger, \n",
    "    $$ \n",
    "    where the operators $\\{K_1, K_2, \\dots, K_m\\}$ are called Kraus operators and must obey \n",
    "    $\\sum_{j=1}^m K_j^\\dagger K_j = I$ to conserve the trace of $\\rho$. \n",
    "    Maps expressed in the above form are called Kraus maps. It can be shown that every physical map on a finite\n",
    "    dimensional quantum system can be represented as a Kraus map, though this representation is not generally unique.\n",
    "    [You can find more information about quantum operations here](https://en.wikipedia.org/wiki/Quantum_operation#Kraus_operators)\n",
    "  \n",
    "In a way, coherent errors are *in principle* amendable by more precisely calibrated control. Incoherent errors are more tricky."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why do incoherent errors happen?\n",
    "When a quantum system (e.g., the qubits on a quantum processor) is not perfectly isolated from its environment it generally co-evolves with the degrees of freedom it couples to. The implication is that while the total time evolution of system and environment can be assumed to be unitary, restriction to the system state generally is not.\n",
    "\n",
    "**Let's throw some math at this for clarity:**\n",
    "Let our total Hilbert space be given by the tensor product of system and environment Hilbert spaces:\n",
    "$\\mathcal{H} = \\mathcal{H}_S \\otimes \\mathcal{H}_E$.\n",
    "Our system \"not being perfectly isolated\" must be translated to the statement that the global Hamiltonian contains a contribution that couples the system and environment:\n",
    "$$\n",
    "H = H_S \\otimes I + I \\otimes H_E + V\n",
    "$$\n",
    "where $V$ non-trivally acts on both the system and the environment.\n",
    "Consequently, even if we started in an initial state that factorized over system and environment $\\ket{\\psi}_{S,0}\\otimes \\ket{\\psi}_{E,0}$\n",
    "if everything evolves by the Schrödinger equation\n",
    "$$\n",
    "\\ket{\\psi_t} = e^{-i \\frac{Ht}{\\hbar}} \\left(\\ket{\\psi}_{S,0}\\otimes \\ket{\\psi}_{E,0}\\right)\n",
    "$$\n",
    "the final state will generally not admit such a factorization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A toy model\n",
    "**In this (somewhat technical) section we show how environment interaction can corrupt an identity gate and derive its Kraus map.**\n",
    "For simplicity, let us assume that we are in a reference frame in which both the system and environment Hamiltonian's vanish $H_S = 0, H_E = 0$ and where the cross-coupling is small even when multiplied by the duration of the time evolution $\\|\\frac{tV}{\\hbar}\\|^2 \\sim \\epsilon \\ll 1$ (any operator norm $\\|\\cdot\\|$ will do here).\n",
    "Let us further assume that $V = \\sqrt{\\epsilon} V_S \\otimes V_E$ (the more general case is given by a sum of such terms) and that \n",
    "the initial environment state satisfies $\\bra{\\psi}_{E,0} V_E\\ket{\\psi}_{E,0} = 0$. This turns out to be a very reasonable assumption in practice but a more thorough discussion exceeds our scope.\n",
    "\n",
    "Then the joint system + environment state $\\rho = \\rho_{S,0} \\otimes \\rho_{E,0}$ (now written as a density matrix) evolves as\n",
    "$$\n",
    "\\rho \\mapsto \\rho' :=  e^{-i \\frac{Vt}{\\hbar}} \\rho e^{+i \\frac{Vt}{\\hbar}}\n",
    "$$\n",
    "Using the Baker-Campbell-Hausdorff theorem we can expand this to second order in $\\epsilon$\n",
    "$$\n",
    "\\rho' = \\rho - \\frac{it}{\\hbar} [V, \\rho] - \\frac{t^2}{2\\hbar^2} [V, [V, \\rho]] + O(\\epsilon^{3/2})\n",
    "$$\n",
    "We can insert the initially factorizable state $\\rho = \\rho_{S,0} \\otimes \\rho_{E,0}$ and trace over the environmental degrees of freedom to obtain\n",
    "\\begin{align}\n",
    "\\rho_S' := \\tr_E \\rho' & = \\rho_{S,0}  \\underbrace{\\tr \\rho_{E,0}}_{1} - \\frac{i\\sqrt{\\epsilon} t}{\\hbar} \\underbrace{\\left[ V_S \\rho_{S,0} \\underbrace{\\tr V_E\\rho_{E,0}}_{\\bra{\\psi}_{E,0} V_E\\ket{\\psi}_{E,0} = 0} - \\rho_{S,0}V_S  \\underbrace{\\tr \\rho_{E,0}V_E}_{\\bra{\\psi}_{E,0} V_E\\ket{\\psi}_{E,0} = 0} \\right]}_0 \\\\\n",
    "& \\qquad - \\frac{\\epsilon t^2}{2\\hbar^2} \\left[ V_S^2\\rho_{S,0}\\tr V_E^2 \\rho_{E,0} + \\rho_{S,0} V_S^2 \\tr \\rho_{E,0}V_E^2 - 2 V_S\\rho_{S,0}V_S\\tr V_E \\rho_{E,0}V_E\\right] \\\\\n",
    "& = \\rho_{S,0} - \\frac{\\gamma}{2} \\left[ V_S^2\\rho_{S,0} + \\rho_{S,0} V_S^2  - 2 V_S\\rho_{S,0}V_S\\right]\n",
    "\\end{align}\n",
    "where the coefficient in front of the second part is by our initial assumption very small $\\gamma := \\frac{\\epsilon t^2}{2\\hbar^2}\\tr V_E^2 \\rho_{E,0} \\ll 1$.\n",
    "This evolution happens to be approximately equal to a Kraus map with operators $K_1 := I - \\frac{\\gamma}{2} V_S^2, K_2:= \\sqrt{\\gamma} V_S$:\n",
    "\\begin{align}\n",
    "\\rho_S \\to \\rho_S' &= K_1\\rho K_1^\\dagger + K_2\\rho K_2^\\dagger\n",
    " = \\rho - \\frac{\\gamma}{2}\\left[ V_S^2 \\rho + \\rho V_S^2\\right] + \\gamma V_S\\rho_S V_S + O(\\gamma^2)\n",
    "\\end{align}\n",
    "This agrees to $O(\\epsilon^{3/2})$ with the result of our derivation above. This type of derivation can be extended to many other cases with little complication and a very similar argument is used to derive the [Lindblad master equation](https://en.wikipedia.org/wiki/Lindblad_equation)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Support for noisy gates on the Rigetti QVM\n",
    "\n",
    "As of today, users of our Forest API can annotate their QUIL programs by certain pragma statements that inform the QVM that a particular gate on specific target qubits should be replaced by an imperfect realization given by a Kraus map.\n",
    "\n",
    "## But the QVM propagates *pure states*: How does it simulate noisy gates?\n",
    "It does so by yielding the correct outcomes **in the average over many executions of the QUIL program**:\n",
    "When the noisy version of a gate should be applied the QVM makes a random choice which Kraus operator is applied to the current state with a probability that ensures that the average over many executions is equivalent to the Kraus map.\n",
    "In particular, a particular Kraus operator $K_j$ is applied to $\\ket{\\psi}_S$\n",
    "$$\n",
    "\\ket{\\psi'}_S = \\frac{1}{\\sqrt{p_j}} K_j \\ket{\\psi}_S\n",
    "$$\n",
    "with probability $p_j:= \\bra{\\psi}_S K_j^\\dagger K_j \\ket{\\psi}_S$.\n",
    "In the average over many execution $N \\gg 1$ we therefore find that\n",
    "\\begin{align}\n",
    "\\overline{\\rho_S'} & = \\frac{1}{N} \\sum_{n=1}^N \\ket{\\psi'_n}_S\\bra{\\psi'_n}_S \\\\\n",
    "& = \\frac{1}{N} \\sum_{n=1}^N p_{j_n}^{-1}K_{j_n}\\ket{\\psi'}_S \\bra{\\psi'}_SK_{j_n}^\\dagger\n",
    "\\end{align}\n",
    "where $j_n$ is the chosen Kraus operator label in the $n$-th trial.\n",
    "This is clearly a Kraus map itself! And we can group identical terms and rewrite it as\n",
    "\\begin{align}\n",
    "\\overline{\\rho_S'} & = \n",
    "  \\sum_{\\ell=1}^n \\frac{N_\\ell}{N}  p_{\\ell}^{-1}K_{\\ell}\\ket{\\psi'}_S \\bra{\\psi'}_SK_{\\ell}^\\dagger\n",
    "\\end{align}\n",
    "where $N_{\\ell}$ is the number of times that Kraus operator label $\\ell$ was selected.\n",
    "For large enough $N$ we know that $N_{\\ell} \\approx N p_\\ell$ and therefore\n",
    "\\begin{align}\n",
    "\\overline{\\rho_S'} \\approx \\sum_{\\ell=1}^n K_{\\ell}\\ket{\\psi'}_S \\bra{\\psi'}_SK_{\\ell}^\\dagger\n",
    "\\end{align}\n",
    "which proves our claim.\n",
    "**The consequence is that noisy gate simulations must generally be repeated many times to obtain representative results**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How do I get started?\n",
    "\n",
    "1. Come up with a good model for your noise. We will provide some examples below and may add more such \n",
    "    examples to our public repositories over time. Alternatively, you can characterize the gate under \n",
    "    consideration using [Quantum Process Tomography](https://arxiv.org/abs/1202.5344) or \n",
    "    [Gate Set Tomography](http://www.pygsti.info/) and use the resulting process matrices to obtain a \n",
    "    very accurate noise model for a particular QPU.\n",
    "2. Define your Kraus operators as a list of numpy arrays `kraus_ops = [K1, K2, ..., Km]`.\n",
    "3. For your QUIL program `p`, call:\n",
    "    ```\n",
    "    p.define_noisy_gate(\"MY_NOISY_GATE\", [q1, q2], kraus_ops)\n",
    "    ```\n",
    "    where you should replace `MY_NOISY_GATE` with the gate of interest and `q1, q2` the indices of the qubits.\n",
    "    \n",
    "**Scroll down for some examples!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import binom\n",
    "import matplotlib.colors as colors\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyquil import Program, get_qc\n",
    "from pyquil.gates import CZ, H, I, X, MEASURE\n",
    "from scipy.linalg import expm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We could ask for \"2q-noisy-qvm\" but we will be specifying\n",
    "# our noise model as PRAGMAs on the Program itself.\n",
    "qc = get_qc('2q-qvm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 1: Amplitude damping\n",
    "\n",
    "Amplitude damping channels are imperfect identity maps with Kraus operators\n",
    "$$\n",
    "K_1 = \\begin{pmatrix} \n",
    "1 & 0 \\\\\n",
    "0 & \\sqrt{1-p}\n",
    "\\end{pmatrix} \\\\\n",
    "K_2 = \\begin{pmatrix} \n",
    "0 & \\sqrt{p} \\\\\n",
    "0 & 0\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "where $p$ is the probability that a qubit in the $\\ket{1}$ state decays to the $\\ket{0}$ state.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def damping_channel(damp_prob=.1):\n",
    "    \"\"\"\n",
    "    Generate the Kraus operators corresponding to an amplitude damping \n",
    "    noise channel.\n",
    "    \n",
    "    :params float damp_prob: The one-step damping probability.\n",
    "    :return: A list [k1, k2] of the Kraus operators that parametrize the map.\n",
    "    :rtype: list\n",
    "    \"\"\"\n",
    "    damping_op = np.sqrt(damp_prob) * np.array([[0, 1],\n",
    "                                                [0, 0]])\n",
    "    \n",
    "    residual_kraus = np.diag([1, np.sqrt(1-damp_prob)])\n",
    "    return [residual_kraus, damping_op]\n",
    "\n",
    "def append_kraus_to_gate(kraus_ops, g):\n",
    "    \"\"\"\n",
    "    Follow a gate `g` by a Kraus map described by `kraus_ops`.\n",
    "    \n",
    "    :param list kraus_ops: The Kraus operators.\n",
    "    :param numpy.ndarray g: The unitary gate.\n",
    "    :return: A list of transformed Kraus operators.\n",
    "    \"\"\"\n",
    "    return [kj.dot(g) for kj in kraus_ops]\n",
    "\n",
    "\n",
    "def append_damping_to_gate(gate, damp_prob=.1):\n",
    "    \"\"\"\n",
    "    Generate the Kraus operators corresponding to a given unitary \n",
    "    single qubit gate followed by an amplitude damping noise channel.\n",
    "    \n",
    "    :params np.ndarray|list gate: The 2x2 unitary gate matrix.\n",
    "    :params float damp_prob: The one-step damping probability.\n",
    "    :return: A list [k1, k2] of the Kraus operators that parametrize the map.\n",
    "    :rtype: list\n",
    "    \"\"\"\n",
    "    return append_kraus_to_gate(damping_channel(damp_prob), gate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# single step damping probability\n",
    "damping_per_I = 0.02\n",
    "\n",
    "# number of program executions\n",
    "trials = 200\n",
    "\n",
    "results_damping = []\n",
    "lengths = np.arange(0, 201, 10, dtype=int)\n",
    "for jj, num_I in enumerate(lengths):\n",
    "    \n",
    "    print(\"\\r{}/{}, \".format(jj, len(lengths)), end=\"\")\n",
    "\n",
    "    \n",
    "    p = Program(X(0))\n",
    "    # want increasing number of I-gates\n",
    "    p.inst([I(0) for _ in range(num_I)])\n",
    "    p.inst(MEASURE(0, 0))\n",
    "    \n",
    "    # overload identity I on qc 0\n",
    "    p.define_noisy_gate(\"I\", [0], append_damping_to_gate(np.eye(2), damping_per_I))\n",
    "    p.wrap_in_numshots_loop(trials)\n",
    "    qc.qam.random_seed = int(num_I)\n",
    "    res = qc.run(p)\n",
    "    results_damping.append([np.mean(res), np.std(res) / np.sqrt(trials)])\n",
    "    \n",
    "results_damping = np.array(results_damping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dense_lengths = np.arange(0, lengths.max()+1, .2)\n",
    "survival_probs = (1-damping_per_I)**dense_lengths\n",
    "logpmf = binom.logpmf(np.arange(trials+1)[np.newaxis, :], trials, survival_probs[:, np.newaxis])/np.log(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DARK_TEAL = '#48737F'\n",
    "FUSCHIA = \"#D6619E\"\n",
    "BEIGE = '#EAE8C6'\n",
    "cm = colors.LinearSegmentedColormap.from_list('anglemap', [\"white\", FUSCHIA, BEIGE], N=256, gamma=1.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(14, 6))\n",
    "plt.pcolor(dense_lengths, np.arange(trials+1)/trials, logpmf.T, cmap=cm, vmin=-4, vmax=logpmf.max())\n",
    "plt.plot(dense_lengths, survival_probs, c=BEIGE, label=\"Expected mean\")\n",
    "plt.errorbar(lengths, results_damping[:,0], yerr=2*results_damping[:,1], c=DARK_TEAL, \n",
    "             label=r\"noisy qvm, errorbars $ = \\pm 2\\hat{\\sigma}$\", marker=\"o\")\n",
    "cb = plt.colorbar()\n",
    "cb.set_label(r\"$\\log_{10} \\mathrm{Pr}(n_1; n_{\\rm trials}, p_{\\rm survival}(t))$\", size=20)\n",
    "\n",
    "plt.title(\"Amplitude damping model of a single qubit\", size=20)\n",
    "plt.xlabel(r\"Time $t$ [arb. units]\", size=14)\n",
    "plt.ylabel(r\"$n_1/n_{\\rm trials}$\", size=14)\n",
    "plt.legend(loc=\"best\", fontsize=18)\n",
    "plt.xlim(*lengths[[0, -1]])\n",
    "plt.ylim(0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example 2: dephased CZ-gate\n",
    "\n",
    "Dephasing is usually characterized through a qubit's $T_2$ time. \n",
    "For a single qubit the dephasing Kraus operators are\n",
    "$$\n",
    "K_1(p) = \\sqrt{1-p} I_2 \\\\\n",
    "K_2(p) = \\sqrt{p} \\sigma_Z\n",
    "$$\n",
    "where $p = 1 - \\exp(-T_2/T_{\\rm gate})$ is the probability that the qubit is dephased over the time interval of interest, $I_2$ is the $2\\times 2$-identity matrix and $\\sigma_Z$ is the Pauli-Z operator.\n",
    "\n",
    "For two qubits, we must construct a Kraus map that has *four* different outcomes:\n",
    "\n",
    "1. No dephasing\n",
    "2. Qubit 1 dephases\n",
    "3. Qubit 2 dephases\n",
    "4. Both dephase"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Kraus operators for this are given by\n",
    "\\begin{align}\n",
    "K'_1(p,q) = K_1(p)\\otimes K_1(q) \\\\\n",
    "K'_2(p,q) = K_2(p)\\otimes K_1(q) \\\\\n",
    "K'_3(p,q) = K_1(p)\\otimes K_2(q) \\\\\n",
    "K'_4(p,q) = K_2(p)\\otimes K_2(q) \n",
    "\\end{align}\n",
    "where we assumed a dephasing probability $p$ for the first qubit and $q$ for the second.\n",
    "\n",
    "Dephasing is a *diagonal* error channel and the CZ gate is also diagonal, therefore we can get the combined map of dephasing and the CZ gate simply by composing $U_{\\rm CZ}$ the unitary representation of CZ with each Kraus operator\n",
    "\\begin{align}\n",
    "K^{\\rm CZ}_1(p,q) = K_1(p)\\otimes K_1(q)U_{\\rm CZ} \\\\\n",
    "K^{\\rm CZ}_2(p,q) = K_2(p)\\otimes K_1(q)U_{\\rm CZ} \\\\\n",
    "K^{\\rm CZ}_3(p,q) = K_1(p)\\otimes K_2(q)U_{\\rm CZ} \\\\\n",
    "K^{\\rm CZ}_4(p,q) = K_2(p)\\otimes K_2(q)U_{\\rm CZ} \n",
    "\\end{align}\n",
    "\n",
    "**Note that this is not always accurate, because a CZ gate is often achieved through non-diagonal interaction Hamiltonians! However, for sufficiently small dephasing probabilities it should always provide a good starting point.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dephasing_kraus_map(p=.1):\n",
    "    \"\"\"\n",
    "    Generate the Kraus operators corresponding to a dephasing channel.\n",
    "\n",
    "    :params float p: The one-step dephasing probability.\n",
    "    :return: A list [k1, k2] of the Kraus operators that parametrize the map.\n",
    "    :rtype: list\n",
    "    \"\"\"\n",
    "    return [np.sqrt(1-p)*np.eye(2), np.sqrt(p)*np.diag([1, -1])]\n",
    "\n",
    "def tensor_kraus_maps(k1, k2):\n",
    "    \"\"\"\n",
    "    Generate the Kraus map corresponding to the composition\n",
    "    of two maps on different qubits.\n",
    "    \n",
    "    :param list k1: The Kraus operators for the first qubit.\n",
    "    :param list k2: The Kraus operators for the second qubit.\n",
    "    :return: A list of tensored Kraus operators.\n",
    "    \"\"\"\n",
    "    return [np.kron(k1j, k2l) for k1j in k1 for k2l in k2]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# single step damping probabilities\n",
    "ps = np.linspace(.001, .5, 200)\n",
    "\n",
    "# number of program executions\n",
    "trials = 500\n",
    "\n",
    "results = []\n",
    "\n",
    "for jj, p in enumerate(ps):\n",
    "\n",
    "    corrupted_CZ = append_kraus_to_gate(\n",
    "    tensor_kraus_maps(\n",
    "        dephasing_kraus_map(p),\n",
    "        dephasing_kraus_map(p)\n",
    "    ), \n",
    "    np.diag([1, 1, 1, -1]))\n",
    "\n",
    "    \n",
    "    print(\"\\r{}/{}, \".format(jj, len(ps)), end=\"\")\n",
    "    \n",
    "    # make Bell-state\n",
    "    p = Program(H(0), H(1), CZ(0,1), H(1))\n",
    "    \n",
    "    p.inst(MEASURE(0, 0))\n",
    "    p.inst(MEASURE(1, 1))\n",
    "    \n",
    "    # overload identity I on qc 0\n",
    "    p.define_noisy_gate(\"CZ\", [0, 1], corrupted_CZ)\n",
    "    p.wrap_in_numshots_loop(trials)\n",
    "    qc.qam.random_seed = jj\n",
    "    res = qc.run(p)\n",
    "    results.append(res)\n",
    "    \n",
    "results = np.array(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Z1s = (2*results[:,:,0]-1.)\n",
    "Z2s = (2*results[:,:,1]-1.)\n",
    "Z1Z2s = Z1s * Z2s\n",
    "\n",
    "Z1m = np.mean(Z1s, axis=1)\n",
    "Z2m = np.mean(Z2s, axis=1)\n",
    "Z1Z2m = np.mean(Z1Z2s, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(14, 6))\n",
    "plt.axhline(y=1.0, color=FUSCHIA, alpha=.5, label=\"Bell state\")\n",
    "\n",
    "plt.plot(ps, Z1Z2m, \"x\", c=FUSCHIA, label=r\"$\\overline{Z_1 Z_2}$\")\n",
    "plt.plot(ps, 1-2*ps, \"--\", c=FUSCHIA, label=r\"$\\langle Z_1 Z_2\\rangle_{\\rm theory}$\")\n",
    "\n",
    "plt.plot(ps, Z1m, \"o\", c=DARK_TEAL, label=r\"$\\overline{Z}_1$\")\n",
    "plt.plot(ps, 0*ps, \"--\", c=DARK_TEAL, label=r\"$\\langle Z_1\\rangle_{\\rm theory}$\")\n",
    "\n",
    "plt.plot(ps, Z2m, \"d\", c=\"k\", label=r\"$\\overline{Z}_2$\")\n",
    "plt.plot(ps, 0*ps, \"--\", c=\"k\", label=r\"$\\langle Z_2\\rangle_{\\rm theory}$\")\n",
    "\n",
    "plt.xlabel(r\"Dephasing probability $p$\", size=18)\n",
    "plt.ylabel(r\"$Z$-moment\", size=18)\n",
    "plt.title(r\"$Z$-moments for a Bell-state prepared with dephased CZ\", size=18)\n",
    "plt.xlim(0, .5)\n",
    "plt.legend(fontsize=18)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
